// Stub.cpp : Defines the entry point for the application.
//

#include "stdafx.h"
#include "Stub.h"

//static PSTUB_DATA pStubData;

int StubEP()
{
    PSTUB_DATA pStubData;

    // get address of STUB_DATA
    // we know that it's located straight before our code.
    __asm{
        call __getmyaddr
__getmyaddr:
        pop eax
        and eax, 0xFFFFF000 // We know that address of structure is the address of section.
                            // section addr aligned to 0x1000
        mov pStubData, eax
    }

    PPEB Peb;
    __asm {
        mov eax, FS:[0x30];
        mov Peb, eax
        nop
    }

    pStubData->dwOriginalEP += Peb->ImageBaseAddress;


	StartOriginalPE( pStubData );

	return 0;
}

void StartOriginalPE( PSTUB_DATA pStubData )
{
	LPVOID pOriginalEP = reinterpret_cast<LPVOID>( pStubData->dwOriginalEP );
	__asm
	{
		jmp pOriginalEP
	}
}

ULONG_PTR GetKernelBase()
{
    ULONG_PTR hKernelBase = 0;
    __asm
    {
        cld
        mov edx, dword ptr FS:[0x30]
        mov edx, dword ptr [edx+0x0C]
        mov edx, dword ptr [edx+0x14]
next_mod:
        mov esi, dword ptr [edx+0x28]
        push 24
        pop ecx
        xor edi, edi
loop_modname:
        xor eax, eax
        lodsb
        cmp al, 'a'
        jl not_lowercase
        sub al, 0x20
not_lowercase:
        ror edi, 13
        add edi, eax
        loop loop_modname
        cmp edi, 0x6A4ABC5B
        mov ebx, dword ptr [edx+0x10]
        mov edx, dword ptr [edx]
        jne next_mod
        mov hKernelBase, ebx
    }

    return hKernelBase;
}

void FillApisTable( PSTUB_DATA pStubData )
{
    ULONG_PTR hKernelBase = GetKernelBase();
    pStubData->hKernel32 = hKernelBase;

    pStubData->pLoadLibrary += hKernelBase;
    pStubData->pGetProcAddress += hKernelBase;
    pStubData->pExitThread += hKernelBase;
}

void ProcessImports( PSTUB_DATA pStubData )
{
    PIMAGE_NT_HEADERS32 pNtHeaders = reinterpret_cast<PIMAGE_NT_HEADERS32>( \
        reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) + pStubData->pImageBase->e_lfanew );

    PIMAGE_OPTIONAL_HEADER32 pOptHead = &pNtHeaders->OptionalHeader;

    // get import data dir
    PIMAGE_DATA_DIRECTORY pImportDir = &pOptHead->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

    PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor = \
        reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>( pImportDir->VirtualAddress + \
        reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) );

    while ( pImportDescriptor->Name != 0 )
    {
        char* szModuleName = reinterpret_cast<char*>(
            reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) + pImportDescriptor->Name );

        ULONG_PTR hModule = reinterpret_cast<PLoadLibraryA>( pStubData->pLoadLibrary )( szModuleName );

        PDWORD pFirstThunk = reinterpret_cast<PDWORD>( pImportDescriptor->FirstThunk +
            reinterpret_cast<DWORD>( pStubData->pImageBase ) );
        PIMAGE_THUNK_DATA pThunksChain = NULL;

        if ( !( pThunksChain = reinterpret_cast<PIMAGE_THUNK_DATA>( pImportDescriptor->OriginalFirstThunk + 
            reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) ) ) )
        {
            pThunksChain = reinterpret_cast<PIMAGE_THUNK_DATA>( pFirstThunk );
        }

        // load all imported apis from definite dll
        while ( pThunksChain->u1.AddressOfData )
        {
            // by ordinal
            if ( pThunksChain->u1.AddressOfData & IMAGE_ORDINAL_FLAGS )
            {
                *pFirstThunk = reinterpret_cast<PGetProcAddress>( pStubData->pGetProcAddress )
                    ( hModule, reinterpret_cast<char*>( MAKEINTRESOURCE( LOWORD ( pThunksChain->u1.Ordinal ) ) ) );
            }
            else // by name
            {
                char* szFunctionName = reinterpret_cast<char*>(
                    reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) + pThunksChain->u1.AddressOfData + 2 );

                *pFirstThunk = reinterpret_cast<PGetProcAddress>( pStubData->pGetProcAddress )
                    ( hModule, szFunctionName );
            }

            ++pThunksChain;
            ++pFirstThunk;
        }

        pImportDescriptor++;
    }
}

void ApplyFixups( PSTUB_DATA pStubData )
{
    PIMAGE_NT_HEADERS32 pNtHeaders = reinterpret_cast<PIMAGE_NT_HEADERS32>( \
        reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) + pStubData->pImageBase->e_lfanew );

    PIMAGE_OPTIONAL_HEADER32 pOptHead = &pNtHeaders->OptionalHeader;

    // get relocs dir
    PIMAGE_DATA_DIRECTORY pRelocsDir = &pOptHead->DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if ( pRelocsDir->Size == 0 )
    {
        // cmon. No relocs found. eek.
        return;
    }

    PIMAGE_BASE_RELOCATION pReloc, pBaseRelocs;

    pReloc = ( pBaseRelocs = \
        reinterpret_cast<PIMAGE_BASE_RELOCATION>( pRelocsDir->VirtualAddress + \
        reinterpret_cast<ULONG_PTR>( pStubData->pImageBase ) ) );

    DWORD dwRelocCount = 0;
    WORD* pwRelocInfo = 0;

    while ( reinterpret_cast<BYTE*>( pReloc ) - reinterpret_cast<BYTE*>( pBaseRelocs ) < pRelocsDir->Size )
    {
        dwRelocCount = ( pReloc->SizeOfBlock - sizeof(*pReloc) ) / sizeof(WORD);
        pwRelocInfo  = reinterpret_cast<WORD*>( reinterpret_cast<BYTE*>( pReloc ) + sizeof(*pReloc) );

        for ( DWORD dwRelocIndex = 0; dwRelocIndex < dwRelocCount; ++ dwRelocIndex, pwRelocInfo++ )
        {
            DWORD virtualFixup = ( (*pwRelocInfo) & 0xFFF ) + pReloc->VirtualAddress;

            switch ( *pwRelocInfo >> 12 ) 
            {
            case IMAGE_REL_BASED_DIR64:
                *reinterpret_cast<__int64*>( 
                    reinterpret_cast<ULONGLONG>( pStubData->pImageBase ) + virtualFixup ) += 
                    pStubData->dwImageBaseDelta;
                break;
            case IMAGE_REL_BASED_HIGHLOW:
                *reinterpret_cast<__int32*>( 
                    reinterpret_cast<DWORD>( pStubData->pImageBase ) + virtualFixup )+= 
                    pStubData->dwImageBaseDelta;
                break;
            default:
                break;
            }

        }

        pReloc = reinterpret_cast<PIMAGE_BASE_RELOCATION>( pwRelocInfo );
    }

}